﻿namespace Ngs.Expressions
{
    #region Using Directives
    // Copyright (c) Microsoft Corporation.  All rights reserved.
    // This source code is made available under the terms of the Microsoft Public License (MS-PL)

    using System;
    using System.Collections.Generic;
    using System.Collections.ObjectModel;
    using System.Linq;
    using System.Linq.Expressions;
    using System.Text;
    using System.Reflection;

    #endregion
    /// <summary>
    /// Replaces references to one specific instance of an expression node with another node
    /// </summary>
    public class ExpressionReplacer : ExpressionVisitor
    {
        Expression searchFor;
        Expression replaceWith;

        private ExpressionReplacer(Expression searchFor, Expression replaceWith)
        {
            this.searchFor = searchFor;
            this.replaceWith = replaceWith;
        }

        public static Expression Replace(Expression expression, Expression searchFor, Expression replaceWith)
        {
            return new ExpressionReplacer(searchFor, replaceWith).Visit(expression);
        }

        public static Expression ReplaceAll(Expression expression, Expression[] searchFor, Expression[] replaceWith)
        {
            for (int i = 0, n = searchFor.Length; i < n; i++)
            {
                expression = Replace(expression, searchFor[i], replaceWith[i]);
            }
            return expression;
        }

        /// <summary>This lets you normalize parameter expressions throughout an expression tree.
        /// Assumes all parameters are using the same name.</summary>
        public static Expression ReplaceParameters(Expression expression, ParameterExpression param)
        {
            return ReplaceAll(expression, new[] { param }, new[] { param });
        }

        protected override Expression Visit(Expression exp)
        {
            if (exp == this.searchFor)
            {
                return this.replaceWith;
            }
            return base.Visit(exp);
        }

        protected override Expression VisitLambda(LambdaExpression lambda)
        {
            if (searchFor.NodeType == ExpressionType.Parameter)
            {
                ParameterExpression parameter = (ParameterExpression)replaceWith;
                Expression body = Visit(lambda.Body);
                return Expression.Lambda(body, new[] { parameter });
            }
            return base.VisitLambda(lambda);
        }

        protected override Expression VisitMemberAccess(MemberExpression m)
        {
            if (searchFor.NodeType == ExpressionType.Parameter)
            {
                MemberTypes memberType = m.Member.MemberType;
                if (memberType == MemberTypes.Property || memberType == MemberTypes.Field)
                {
                    return Expression.PropertyOrField(replaceWith, m.Member.Name);
                }
            }

            return base.VisitMemberAccess(m);
        }

    }
}

